__doc__ =   """ Contains the MultiEntity class a special Entity class to hold multiple instances
                of other entities. All Entities apart from Atom can hold others and inherit from
                the MultiEntity, use classes in the Entities or Holders modules.
            """
# Classes inheriting from MultiEntity have to provide some attributes during init:
# self.level = char (inside the hierarchy)

from Entity import Entity
from ResidueProperties import AA_NAMES
from ResidueProperties import HOH_NAMES
from AtomProperties import AT_ORDER
from AtomProperties import AT_REMOTE
# do not add more dependancies inside zenPDB!
from operator import itemgetter             # used for sorting list of tuples
from operator import (gt, ge, lt, le)       # binary operators from the operator
from operator import (eq, ne, or_, and_, contains, is_, is_not) # module are needed to make
from collections import defaultdict                            # selectors general
from itertools import izip
from copy import (copy, deepcopy)
from numpy import (mean, sum)

ALPHABET = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ_ '

def sort_id_list(id_list, sort_tuple):
    """ Sorts lists of id tuples.
    """
    (hol_loc, str_loc, mod_loc, chn_loc, res_loc, at_loc) = sort_tuple
    # evan a simple id is a tuple, this makes sorting general
    def space_last(ch_id1, ch_id2):             # this is for chain sorting
        if ch_id1 == ' '  and ch_id2 != ' ':
            return 1
        if ch_id2 == ' '  and ch_id1 != ' ':
            return -1
        if ch_id1 == ' '  and ch_id2 == ' ':
            return 0
        return cmp(ch_id1, ch_id2)

    def atom(at_id1, at_id2):
        # hydrogen atoms come last
        is_hydrogen1 = (at_id1[0] == 'H')
        is_hydrogen2 = (at_id2[0] == 'H')
        diff = cmp(is_hydrogen1, is_hydrogen2)

        # back bone come first
        if not diff:
            order1 = AT_ORDER.get(at_id1)
            order2 = AT_ORDER.get(at_id2)
            diff = cmp(order2, order1)

        # (B)eta, (D)elta, (G)amma, .... o(X)t
        if not diff:
            remote1 = AT_REMOTE.get(at_id1[1:2])
            remote2 = AT_REMOTE.get(at_id2[1:2])
            diff = cmp(remote1, remote2)

        # branching comes last
        if not diff:
            diff = cmp(at_id1[2:4], at_id2[2:4])
        return diff

        # SE vs CE - selenium first
        if not diff:
            alpha1 = ALPHABET.index(at_id1[0:1])
            alpha2 = ALPHABET.index(at_id2[0:1])
            diff = cmp(alpha2, alpha1)

    def residue(res_id1, res_id2):
        r1, r2 = 1, 1
        if res_id1 in AA_NAMES: r1 = 2
        if res_id1 in HOH_NAMES: r1 = 0
        if res_id2 in AA_NAMES: r2 = 2
        if res_id2 in HOH_NAMES: r2 = 0
        if r1 is r2:
            return cmp(res_id1, res_id2)
        else:
            return cmp(r2, r1)
    # this assumes that the implementation of sorting is stable.
    # does it work for others then cPython.
    if res_loc or res_loc is 0:
        id_list.sort(key =itemgetter(res_loc), cmp= lambda x, y: residue(x[0], y[0]))  # by res_name
    if at_loc or at_loc is 0:
        id_list.sort(key =itemgetter(at_loc), cmp= lambda x, y: space_last(x[1], y[1]))  # by alt_loc
    if at_loc or at_loc is 0:
        id_list.sort(key =itemgetter(at_loc), cmp= lambda x, y: atom(x[0], y[0]))  # by at_id
    if res_loc or res_loc is 0:
        id_list.sort(key =itemgetter(res_loc), cmp= lambda x, y: cmp(x[2], y[2]))  # by res_ic
    if res_loc or res_loc is 0:
        id_list.sort(key =itemgetter(res_loc), cmp= lambda x, y: cmp(x[1], y[1]))  # by res_id
    if chn_loc or chn_loc is 0:
        id_list.sort(key =itemgetter(chn_loc), cmp= space_last) # by chain
    if mod_loc or mod_loc is 0:
        id_list.sort(key =itemgetter(mod_loc))             # by model
    if str_loc or str_loc is 0:
        id_list.sort(key =itemgetter(str_loc))             # by structure
    return id_list

def merge(dicts):
    """ Merges multiple dictionaries into a new one.
    """
    master_dict = {}
    for dict_ in dicts:
        master_dict.update(dict_)
    return master_dict

def unique(lists):
    """ Merges multiple iterables into a tuple.
    """
    master_set = set()
    for set_ in lists:
        master_set.update(set_)
    return tuple(sorted(master_set))


class MultiEntity(Entity):
    """ Entity holding entities.
    """
    def __init__(self, long_id, short_id =None, *args):
        self.index = self.hierarchy.index(self.level)   # index corresponding to the hierarchy level
        self.table = dict([(level,{}) for level in self.hierarchy[self.index+1:]])  # empty table
        Entity.__init__(self, long_id, short_id, *args)

    def __repr__(self):
        """ Children-classes should override it to make it meaningfull
        """
        id_ = self.get_id()
        return "<MultiEntity id=%s, holding=%s>" % (id_, len(self))

    def __getstate__(self):
        self.unlink()
        # this is needed for pickling, if the need arises
        # it should also care about __slots__.
        return (self, self.__dict__)

    def __setstate__(self, state):
        self.update(state[0])
        self.__dict__.update(state[1])
        self.link()

    def __iter__(self):
        return self.itervalues()

    def set_sort_tuple(self, sort_tuple =None):
        if sort_tuple:
            self.sort_tuple = sort_tuple
        else:   # making the sort tuple, ugly, uggly, uaughhlly ble
            sort_tuple = [None, None, None, None, None, None]
            key_lenght = len(self.keys()[0])
            stop_i  = self.index + 2                    # next level, open right [)
            start_i = stop_i - key_lenght               # before all nones
            indexes = range(start_i, stop_i)            # Nones to change
            for value, index in enumerate(indexes):
                sort_tuple[index] = value
            self.sort_tuple = sort_tuple

    def get_sort_tuple(self):
        if not hasattr(self, 'sort_tuple'):
           self.set_sort_tuple()
        return self.sort_tuple

    def itervalues(self, unmask =False):
        return (v for v in super(MultiEntity, self).itervalues() if not v.masked or unmask)

    def iteritems(self, unmask =False):
        return ((k,v) for k,v in super(MultiEntity, self).iteritems() if not v.masked or unmask)

    def iterkeys(self, unmask =False):
        return (k for k,v in super(MultiEntity, self).iteritems() if not v.masked or unmask)

    def values(self, *args, **kwargs):
        return list(self.itervalues(*args, **kwargs))

    def items(self, *args, **kwargs):
        return list(self.iteritems(*args, **kwargs))

    def keys(self, *args, **kwargs):
        return list(self.iterkeys(*args, **kwargs))

    def sorted_keys(self, *args, **kwargs):
        list_ = sort_id_list(self.keys(*args, **kwargs), self.get_sort_tuple())
        return list_

    def sorted_values(self, *args, **kwargs):
        values = [self[i] for i in self.sorted_keys(*args, **kwargs)]
        return values

    def sorted_items(self, *args, **kwargs):
        items = [(i, self[i]) for i in self.sorted_keys()]
        return items

    def _set_masked(self, masked):
        """
        """
        if masked:
            # we have to mask children
            for child in self.itervalues():             # only unmasked children
                child._set_masked(True)
                child.set_modified(False, False)
        else:
            # we have to unmask children
            for child in self.itervalues(unmask =True):
                if child.masked:                        # only masked children
                    child._set_masked(False)
                    child.set_modified(False, False)
        self.masked = masked

    def set_masked(self, masked):
        if masked != self.masked:
            self._set_masked(masked)         # mask or unmask
            self.set_modified(True, False)   # set parents as modified

    def set_modified(self, up =True, down =True):
        self.modified = True
        if up and self.parent:
            self.parent.set_modified(True, False)
        if down:
            for child in self.itervalues(unmask =True):
                child.set_modified(False, True)

    def _init_child(self, child):
        child.parent = self
        self[child.get_id()] = child

    def add_child(self, child):
        """ Add a child to the Entity.
        """
        child.set_parent(self)
        child_id = child.get_id()
        self[child_id] = child
        self.set_modified(True, False)

    def del_child(self, child_id):
        """ Remove a child.
        """
        child = self.pop(child_id)
        child.del_parent()
        self.set_modified(True, False)

    def get_children(self, ids =None):
        """ Return a copy of the list of children.
        """
        if ids:
            children = []
            for (id_, child) in self.iteritems():
                if id_ in ids:
                    children.append(child)
        else:
            children = self.values()
        return children

    def link(self):
        """ Recursivly adds parent pointer to children
        """
        for child in self.itervalues(unmask =True):
            child.set_parent(self)
            try:
                child.link()
            except:
                pass

    def unlink(self):
        """ Recursivly deletes parent pointer from children
        """
        for child in self.itervalues(unmask =True):
            child.del_parent()
            try:
                child.unlink()
            except AttributeError:
                pass

    def _set_table(self, entity):
        """ Helper method for self.set_table
        """
        for e in entity.itervalues():
            self.table[e.get_level()].update({e.get_full_id():e})
            self._set_table(e)

    def set_table(self, force =True):
        """ Create a table-dictionary with all children grouped in levels.
        """
        if self.modified or force:
            # a table is accurate as long as the contents of a dictionary do not change.
            self.del_table()
            self._set_table(self)
        self.modified = False

    def del_table(self):
        """ Delete a table-dictionary.
        """
        self.table = dict([(level,{}) for level in self.hierarchy[self.index+1:]])
        self.modified = True

    def get_table(self, level):
        return self.table[level]

    def update_ids(self):
        ids = []
        for (id_, child) in self.iteritems():
            new_id = child.get_id()
            if id_ != new_id:
                ids.append((id_, new_id))
        for (old_id, new_id)  in ids:
            child = self.pop(old_id)
            self.update(((new_id, child),))

    def _data_children(self, attr, xtra =False, function =False, forgiving =True):
        """ Get data from children attributes, methods and xtra dicts.
        """
        if xtra:                                                # looking inside the xtra of children ...
            data = [child.xtra.get(attr) for child in self]     # could get None
        else:                                                   # looking at attributes
            data = []
            for child in self:
                try:
                    if not function:
                        data.append(getattr(child, attr))
                    else:
                        data.append(getattr(child, attr)())
                except AttributeError: #
                    data.append(None)
        if forgiving: # remove Nones
            data = [point for point in data if point is not None]
        return data

    def _data_propagate(self, function, level, *args, **kwargs):
        """ propapagate data from level to calling entity
        """
        if self.index <= self.hierarchy.index(level) - 2:
            for child in self.itervalues():
                child._data_propagate(function, level, *args, **kwargs)
        datas = self._data_children(*args, **kwargs)
        self.xtra[args[0]] = eval(function)(datas)
        return self.xtra[args[0]]

    def _count_children(self, *args, **kwargs):
        data = self._data_children(*args, **kwargs)
        children = defaultdict(int) # by default returns 0
        for d in data:
            children[d] += 1
        return children

    def _freq_children(self, *args, **kwargs):
        children_count = self._count_children(*args, **kwargs)
        lenght = float(len(self))  # it could be len(children_count)?
        for (key_, value_) in children_count.iteritems():
            children_count[key_] = value_/lenght
        return children_count

    def _split_children(self, *args, **kwargs):
        kwargs['forgiving'] = False
        data = self._data_children(*args, **kwargs)
        clusters = defaultdict(dict) # by default returns {}
        for (key, (id_, child)) in izip(data, self.iteritems()):
            clusters[key].update({id_:child})
        return clusters

    def _select_children(self, value, operator, *args, **kwargs):
        """ Generic select children.
        """
        kwargs['forgiving'] = False
        data = self._data_children(*args, **kwargs)
        children = {}
        for (got, (id_, child)) in izip(data, self.iteritems()):
            if eval(operator)(value, got):
                children.update({id_:child})
        return children

    def _ornament_children(self, *args, **kwargs):
        kwargs['forgiving'] = False
        data = self._data_children(*args, **kwargs)
        children = []
        for (got, (id_, child)) in izip(data, self.iteritems()):
            children.append((got, (id_, child)))
        return children

    def _propertydict_children(self, *args, **kwargs):
        kwargs['forgiving'] = False
        data = self._data_children(*args, **kwargs)
        propertydict = {}
        for (got, id_) in izip(data, self.iterkeys()):
            propertydict.update(((id_, got),))
        return propertydict

    def _strip_children(self, *args, **kwargs):
        """ Strip children based on selection criteria.
        """
        children_ids = self._select_children(*args, **kwargs).keys()
        for id_ in children_ids:
            self.del_child(id_)

    def _mask_children(self, *args, **kwargs):
        """ Mask children based on selection criteria.
        """
        children = self._select_children(*args, **kwargs).itervalues()
        for child in children:
            child.set_masked(True) # child.set_modified child.parent.set_modified ...

    def recursive_move(self, origin):
        for child in self.itervalues():
            try:
                child.recursive_move(origin)
            except:
                # Atoms do not have this
                child.move(origin)
                pass
        self.set_coords()

    def recursive_set_coords(self):
        for child in self.itervalues():
            try:
                child.recursive_set_coords()
            except:
                #Atoms do not have this
                pass
        self.set_coords()

    def set_coords(self, *args, **kwargs):
        # select only some children
        if args or kwargs:
            children = self._select_children(*args, **kwargs).values()
        else:
            children = self
        coords = []
        for child in children:
            coords.append(child.get_coords())
        self.coords = mean(coords, axis =0)

    def get_coords(self):
        try:
            return self.coords
        except AttributeError:
            raise AttributeError, 'Entity has coords attribute not set.'

    def dispatch(self, attr, *args, **kwargs):
        for child in self.itervalues():
            getattr(child, attr)(*args, **kwargs)

#EOF